#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
异步任务管理器
提供通用的异步任务处理框架，支持任务队列、状态跟踪和数据库持久化
"""

import queue
import threading
import time
from typing import Optional, Dict, List, Callable, Any, Union
import uuid
import traceback
from datetime import datetime
import json

from models.tasks import Tasks
from models.base import db
from settings.config import logger


class TaskManager:
    """
    通用异步任务管理器
    提供任务队列、状态跟踪和数据库持久化功能
    """
    
    def __init__(self, name: str = "default"):
        """
        初始化任务管理器
        
        Args:
            name: 任务管理器名称，用于区分不同的任务管理器实例
        """
        self.name = name
        self.task_queue = queue.Queue()
        self._worker_thread = None
        self._is_running = False
        self._task_handlers = {}
    
    def register_task_handler(self, task_type: str, handler: Callable):
        """
        注册任务处理器
        
        Args:
            task_type: 任务类型
            handler: 任务处理函数，接收task_id和params参数
        """
        self._task_handlers[task_type] = handler
        logger.info(f"任务处理器已注册: {task_type}")
    
    def submit_task(self, task_type: str, params: Dict = None, account_id: str = None) -> Dict:
        """
        提交异步任务
        
        Args:
            task_type: 任务类型，必须已注册对应的处理器
            params: 任务参数
            account_id: 关联的账号ID
            
        Returns:
            dict: 任务信息，包含task_id
        """
        # 检查任务类型是否已注册处理器
        if task_type not in self._task_handlers:
            return {
                'success': False,
                'message': f'未注册的任务类型: {task_type}',
                'task_id': None
            }
        
        # 生成任务ID
        task_id = f"{task_type}_{int(time.time())}_{str(uuid.uuid4())[:8]}"
        
        # 创建任务记录
        task = Tasks(
            task_id=task_id,
            task_type=task_type,
            status='pending',
            message='任务已提交，等待执行',
            account_id=account_id,
            params=params or {}
        )
        
        try:
            # 保存到数据库
            db.session.add(task)
            db.session.commit()
            
            # 将任务加入队列
            self.task_queue.put({
                'task_id': task_id,
                'task_type': task_type,
                'params': params or {},
                'account_id': account_id
            })
            
            # 启动工作线程（如果还没启动）
            if not self._is_running:
                self._start_worker()
            
            logger.info(f"任务已提交: {task_id}, 类型: {task_type}")
            
            return {
                'success': True,
                'message': '任务已提交',
                'task_id': task_id
            }
            
        except Exception as e:
            db.session.rollback()
            logger.error(f"提交任务失败: {str(e)}")
            return {
                'success': False,
                'message': f'提交任务失败: {str(e)}',
                'task_id': task_id
            }
    
    def get_task_status(self, task_id: str) -> Dict:
        """
        获取任务状态
        
        Args:
            task_id: 任务ID
            
        Returns:
            dict: 任务状态信息
        """
        try:
            task = Tasks.query.filter_by(task_id=task_id).first()
            
            if not task:
                return {
                    'success': False,
                    'message': '任务不存在',
                    'task_id': task_id,
                    'status': 'not_found'
                }
            
            return {
                'success': True,
                'task_id': task_id,
                'status': task.status,
                'message': task.message,
                'progress': task.progress,
                'result': task.result,
                'error': task.error,
                'started_at': task.started_at,
                'finished_at': task.finished_at,
                'created_at': task.created_at.isoformat() if task.created_at else None,
                'updated_at': task.updated_at.isoformat() if task.updated_at else None
            }
            
        except Exception as e:
            logger.error(f"获取任务状态失败: {str(e)}")
            return {
                'success': False,
                'message': f'获取任务状态失败: {str(e)}',
                'task_id': task_id,
                'status': 'error'
            }
    
    def update_task_status(self, task_id: str, status: str, message: str = None, 
                          progress: int = None, result: Dict = None, error: str = None):
        """
        更新任务状态
        
        Args:
            task_id: 任务ID
            status: 任务状态 (pending, running, completed, failed)
            message: 状态消息
            progress: 进度 (0-100)
            result: 任务结果
            error: 错误信息
        """
        try:
            # 检查是否在Flask应用上下文中
            try:
                from flask import has_app_context
                if not has_app_context():
                    logger.warning(f"更新任务状态时没有应用上下文: {task_id}, 尝试创建上下文")
                    try:
                        from api_restructured import create_app
                        app = create_app()
                        with app.app_context():
                            return self._update_task_status_impl(task_id, status, message, progress, result, error)
                    except Exception as e:
                        logger.error(f"创建应用上下文失败: {str(e)}")
                        return False
            except ImportError:
                # 如果不是Flask应用，继续执行
                pass
                
            return self._update_task_status_impl(task_id, status, message, progress, result, error)
            
        except Exception as e:
            logger.error(f"更新任务状态失败: {str(e)}")
            logger.error(traceback.format_exc())
            return False
            
    def _update_task_status_impl(self, task_id: str, status: str, message: str = None, 
                              progress: int = None, result: Dict = None, error: str = None):
        """
        更新任务状态的实际实现
        """
        try:
            task = Tasks.query.filter_by(task_id=task_id).first()
            
            if not task:
                logger.warning(f"更新状态失败: 任务不存在 {task_id}")
                return False
            
            # 更新状态
            task.status = status
            
            if message is not None:
                task.message = message
                
            if progress is not None:
                task.progress = progress
                
            if result is not None:
                task.result = result
                
            if error is not None:
                task.error = error
            
            # 记录开始和结束时间
            current_time = datetime.now().isoformat()
            
            if status == 'running' and not task.started_at:
                task.started_at = current_time
                
            if status in ['completed', 'failed']:
                task.finished_at = current_time
            
            # 保存到数据库
            db.session.commit()
            logger.info(f"任务状态已更新: {task_id}, 状态: {status}")
            return True
            
        except Exception as e:
            db.session.rollback()
            logger.error(f"更新任务状态失败: {str(e)}")
            return False
    
    def _start_worker(self):
        """
        启动工作线程
        """
        if self._worker_thread and self._worker_thread.is_alive():
            return
            
        self._is_running = True
        self._worker_thread = threading.Thread(target=self._worker_loop, daemon=True)
        self._worker_thread.start()
        logger.info(f"工作线程已启动: {self.name}")
    
    def _worker_loop(self):
        """
        工作线程循环
        处理队列中的任务
        """
        # 导入Flask应用
        from flask import current_app
        
        while self._is_running:
            try:
                # 从队列中获取任务（超时1秒）
                task = self.task_queue.get(timeout=1)
                
                # 在Flask应用上下文中处理任务
                if current_app:
                    with current_app.app_context():
                        self._process_task(task)
                else:
                    # 尝试获取应用实例
                    try:
                        from api_restructured import create_app
                        app = create_app()
                        with app.app_context():
                            self._process_task(task)
                    except Exception as e:
                        logger.error(f"无法创建应用上下文: {str(e)}")
                        self._process_task_without_context(task)
                
                self.task_queue.task_done()
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"工作线程错误: {str(e)}")
                logger.error(traceback.format_exc())
    
    def _process_task(self, task: Dict):
        """
        处理任务
        
        Args:
            task: 任务信息
        """
        task_id = task['task_id']
        task_type = task['task_type']
        params = task.get('params', {})
        account_id = task.get('account_id')
        
        # 获取任务处理器
        handler = self._task_handlers.get(task_type)
        
        if not handler:
            self.update_task_status(
                task_id, 
                'failed', 
                f'未找到任务处理器: {task_type}',
                error=f'未找到任务处理器: {task_type}'
            )
            return
        
        try:
            # 更新状态为进行中
            self.update_task_status(task_id, 'running', '任务执行中...')
            
            # 执行任务处理器
            result = handler(task_id=task_id, params=params, account_id=account_id)
            
            # 更新状态为完成
            self.update_task_status(
                task_id, 
                'completed', 
                '任务已完成',
                progress=100,
                result=result
            )
            
        except Exception as e:
            # 获取完整的异常堆栈
            error_stack = traceback.format_exc()
            logger.error(f"任务执行失败: {task_id}, 错误: {str(e)}\n{error_stack}")
            
            # 更新状态为失败
            self.update_task_status(
                task_id, 
                'failed', 
                f'任务执行失败: {str(e)}',
                error=error_stack
            )
    
    def stop(self):
        """
        停止工作线程
        """
        self._is_running = False
        if self._worker_thread:
            self._worker_thread.join(timeout=5)
            logger.info(f"工作线程已停止: {self.name}")


    def _process_task_without_context(self, task: Dict):
        """
        在没有应用上下文的情况下处理任务
        仅记录任务信息，不更新数据库
        
        Args:
            task: 任务信息
        """
        task_id = task['task_id']
        task_type = task['task_type']
        params = task.get('params', {})
        account_id = task.get('account_id')
        
        # 获取任务处理器
        handler = self._task_handlers.get(task_type)
        
        if not handler:
            logger.error(f"未找到任务处理器: {task_type}, 任务ID: {task_id}")
            return
        
        try:
            # 记录任务开始
            logger.info(f"开始执行任务(无上下文): {task_id}, 类型: {task_type}")
            
            # 执行任务处理器
            result = handler(task_id=task_id, params=params, account_id=account_id)
            
            # 记录任务完成
            logger.info(f"任务执行完成(无上下文): {task_id}, 类型: {task_type}")
            
        except Exception as e:
            # 获取完整的异常堆栈
            error_stack = traceback.format_exc()
            logger.error(f"任务执行失败(无上下文): {task_id}, 错误: {str(e)}\n{error_stack}")


# 创建默认任务管理器实例
default_task_manager = TaskManager("default")

# 提供初始化函数，在Flask应用上下文中调用
def init_task_manager(app):
    """
    在Flask应用上下文中初始化任务管理器
    应在应用启动时调用
    
    Args:
        app: Flask应用实例
    """
    # 在应用启动时初始化任务管理器
    with app.app_context():
        logger.info("初始化任务管理器...")
        # 确保任务管理器在应用上下文中启动
        if not default_task_manager._is_running:
            default_task_manager._start_worker()
            
    # 应用关闭时停止任务管理器
    @app.teardown_appcontext
    def shutdown_task_manager(exception=None):
        if default_task_manager._is_running:
            default_task_manager.stop()